# 导入必要的库和模块
import gradio as gr
import numpy as np
import pickle
from PIL import Image
import cv2

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)


def _to_numpy(image):
    """确保输入为 numpy 数组灰度图（uint8）。"""
    if image is None:
        return None
    if isinstance(image, dict):
        # 避免对 numpy 数组使用布尔上下文（image.get(...) or ... 会触发 ValueError）
        if 'image' in image:
            image = image['image']
        elif 'composite' in image:
            image = image['composite']
        else:
            image = image
    if isinstance(image, Image.Image):
        image = np.array(image)
    if not isinstance(image, np.ndarray):
        image = np.array(image)
    # If RGBA or RGB
    if image.ndim == 3:
        if image.shape[2] == 4:
            # use alpha channel as ink
            alpha = image[:, :, 3]
            gray = 255 - alpha
        else:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        gray = image
    return gray.astype(np.uint8)


def preprocess_for_digits(gray: np.ndarray) -> np.ndarray:
    """把任意手写图像转换为与 sklearn digits 数据集一致的特征向量 (1,64).

    流程：灰度->反转（若背景白）->阈值->裁剪->居中->模糊->缩放到8x8->归一化到0..16
    """
    # 自动反转（如果背景偏白）
    if np.mean(gray) > 127:
        gray = 255 - gray

    # 二值化（Otsu）以便更好地找到边界
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # 找到非零像素位置
    coords = np.column_stack(np.where(thresh > 0))
    if coords.size == 0:
        return None

    y0, x0 = coords.min(axis=0)
    y1, x1 = coords.max(axis=0)

    # 添加小边距
    pad = 8
    y0 = max(0, y0 - pad)
    x0 = max(0, x0 - pad)
    y1 = min(gray.shape[0], y1 + pad)
    x1 = min(gray.shape[1], x1 + pad)

    roi = gray[y0:y1, x0:x1]

    # 放到正方形画布上并居中
    h, w = roi.shape
    size = max(h, w)
    square = np.zeros((size, size), dtype=np.uint8)
    y_off = (size - h) // 2
    x_off = (size - w) // 2
    square[y_off:y_off+h, x_off:x_off+w] = roi

    # 计算质心并将其移动到中心
    M = cv2.moments(square)
    if M['m00'] != 0:
        cx = M['m10'] / M['m00']
        cy = M['m01'] / M['m00']
        shift_x = (size / 2) - cx
        shift_y = (size / 2) - cy
        T = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
        square = cv2.warpAffine(square, T, (size, size), borderValue=0)

    # 平滑降噪，帮助 KNN 筛除细小噪点
    square = cv2.GaussianBlur(square, (3, 3), 0)

    # 缩放到 8x8
    small = cv2.resize(square, (8, 8), interpolation=cv2.INTER_AREA)

    # 归一化到 0..16（digits 数据集范围）
    small = small.astype(np.float32)
    # 对比度拉伸到 0..255 再映射到 0..16
    if small.max() > small.min():
        small = (small - small.min()) / (small.max() - small.min()) * 255.0
    small = small / 255.0 * 16.0

    return small.flatten().reshape(1, -1)


def predict_digit_simple(image) -> str:
    """返回单个预测结果（字符串），无调试信息。"""
    try:
        gray = _to_numpy(image)
        if gray is None:
            return ""

        # 尝试更严格的预处理
        features = preprocess_for_digits(gray)

        # 如果严格预处理失败（例如没有检测到笔迹），尝试直接缩放到8x8作为回退
        if features is None:
            # 直接缩放：先确保为 uint8，反转背景为白->黑以匹配训练
            arr = gray.astype(np.uint8)
            if np.mean(arr) > 127:
                arr = 255 - arr
            small = cv2.resize(arr, (8, 8), interpolation=cv2.INTER_AREA).astype(np.float32)
            if small.max() > small.min():
                small = (small - small.min()) / (small.max() - small.min()) * 255.0
            small = small / 255.0 * 16.0
            features = small.flatten().reshape(1, -1)

        pred = knn_model.predict(features)[0]
        return str(int(pred))
    except Exception as e:
        # 打印到控制台便于调试（Gradio UI 不显示）
        import traceback

        traceback.print_exc()
        return ""


def clear_all():
    # 返回值对应 [sketchpad, result_text]
    return None, ""


with gr.Blocks(title="KNN 手写数字识别") as demo:
    gr.Markdown("# KNN 手写数字识别\n在左侧画布写数字，右侧显示预测结果。点击 Clear 同时清空画布与结果。")

    with gr.Row():
        with gr.Column(scale=2):
            # 兼容不同版本的 Gradio：优先尝试使用 Sketchpad 带 canvas_size 参数，失败时退回到 Image 画布
            try:
                sketch = gr.Sketchpad(image_mode="RGB", canvas_size=(280, 280), label="手写区")
            except Exception:
                sketch = gr.Image(source="canvas", tool="sketch", type="pil", label="手写区")
            with gr.Row():
                predict_btn = gr.Button("识别")
                clear_btn = gr.Button("清除")
        with gr.Column(scale=1):
            result = gr.Textbox(label="预测结果", lines=1)

    predict_btn.click(fn=predict_digit_simple, inputs=[sketch], outputs=[result])
    clear_btn.click(fn=clear_all, inputs=None, outputs=[sketch, result])


if __name__ == '__main__':
    demo.launch(share=False)
